
#include "three_nn_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {

constexpr int32_t CORE_NUM = 1;

constexpr int32_t X1_INDEX = 0;
constexpr int32_t X2_INDEX = 1;
constexpr int32_t BYTE_BLOCK = 32;
constexpr int32_t BUF_COUNT = 3;

static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    ThreeNNTilingData tiling;
    uint32_t totalCoreNum = CORE_NUM;
    uint64_t ubSizePlatForm;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);

    const gert::Shape x1Shape = context->GetInputShape(X1_INDEX)->GetStorageShape();
    const gert::Shape x2Shape = context->GetInputShape(X2_INDEX)->GetStorageShape();
    int32_t dTypeSize = 4;
    int32_t perBlockCount = BYTE_BLOCK / dTypeSize;
    int64_t mSize = x2Shape.GetDim(1);
    int64_t singleCount = mSize * 3;

    uint64_t ubMaxProcCount = ubSizePlatForm / BUF_COUNT / 3 / BYTE_BLOCK * perBlockCount * 3;
    context->SetTilingKey(101);
    uint64_t loopTime = 0, tailCount = 0;
    if (singleCount > ubMaxProcCount) {
        context->SetTilingKey(201);
        loopTime = singleCount / ubMaxProcCount;
        tailCount = singleCount % ubMaxProcCount;
    }

    tiling.set_ubMaxProcCount(ubMaxProcCount);
    tiling.set_bSize(x1Shape.GetDim(0));
    tiling.set_nSize(x1Shape.GetDim(1));
    tiling.set_mSize(mSize);
    tiling.set_loopTime(loopTime);
    tiling.set_tailCount(tailCount);

    context->SetBlockDim(totalCoreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    printf(
        "==zf==Tiling Data, ubSizePlatForm:%lu, ubMaxProcCount:%lu, bSize:%lu, nSize:%lu, mSize:%lu, loopTime:%lu, "
        "tailCount:%lu\n",
        ubSizePlatForm, tiling.get_ubMaxProcCount(), tiling.get_bSize(), tiling.get_nSize(), tiling.get_mSize(),
        tiling.get_loopTime(), tiling.get_tailCount());

    // 以下代码是ascendebug必须的
    size_t* workspaces = context->GetWorkspaceSizes(1);
    workspaces[0] = 16 * 1024 * 1024;
    return ge::GRAPH_SUCCESS;

    return ge::GRAPH_SUCCESS;
}
}  // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    return GRAPH_SUCCESS;
}
}  // namespace ge

namespace ops {
class ThreeNN : public OpDef {
public:
    explicit ThreeNN(const char* name) : OpDef(name) {
        this->Input("xyz1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("xyz2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("dist")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(ThreeNN);
}  // namespace ops
